/**
 * @file BMatrix.cpp
 * @brief Computes matrices for membrane and bending calculations
 * 
 */

#include <vector>
#include <array>
#include <cmath>
#include <iostream>

// Define precision type
#ifdef DOUBLE_PRECISION
using Real = double;
#else
using Real = float;
#endif

// Constants and parameters
constexpr int NTSHEL = 8; // Number of shell thickness integration points
constexpr int LGR = 24;   // Dimension size (from Fortran PARAM.h)

// Shape functions data (would normally come from MATELE_COM module)
extern std::array<std::array<Real, 4>, 4> TSHAPS; // Shape function derivatives in ksi
extern std::array<std::array<Real, 4>, 4> TSHAPT; // Shape function derivatives in eta

/**
 * @brief Computes B matrices for membrane and bending calculations
 * 
 * @param elementTypes Vector of element types (ITS1)
 * @param thicknessCoords Thickness coordinates (TSCZ1)
 * @param displacementDerivatives Displacement derivatives (UD)
 * @param b1Matrix Output B1 matrix (LGR x 24 x 3)
 * @param b2Matrix Output B2 matrix (LGR x 24 x 3)
 * @param b11Matrix Output B11 matrix (LGR x 24 x 3)
 * @param b12Matrix Output B12 matrix (LGR x 24 x 3)
 * @param globalMatrix Global B matrix (12 x 4*NTSHEL x 8)
 * @param startIndex Starting index (IBEG)
 * @param gaussPoint Current Gauss point index (IG)
 * @param numElements Number of elements to process (LLT)
 */
void ComputeBMatrix(
    const std::vector<int>& elementTypes,
    const std::vector<std::vector<std::vector<Real>>>& thicknessCoords,
    const std::vector<std::vector<std::vector<Real>>>& displacementDerivatives,
    std::vector<std::vector<std::vector<Real>>>& b1Matrix,
    std::vector<std::vector<std::vector<Real>>>& b2Matrix,
    std::vector<std::vector<std::vector<Real>>>& b11Matrix,
    std::vector<std::vector<std::vector<Real>>>& b12Matrix,
    const std::vector<std::vector<std::vector<Real>>>& globalMatrix,
    int startIndex,
    int gaussPoint,
    int numElements)
{
    // Constants
    const Real Z = -1.0 / std::sqrt(3.0);
    const Real Z2 = Z * Z;
    const int gaussPointHigh = gaussPoint + 4;

    // Extract shape function derivatives for current Gauss point
    const Real DK1 = 0.5 * TSHAPS[0][gaussPoint];
    const Real DE1 = 0.5 * TSHAPT[0][gaussPoint];
    const Real DK2 = 0.5 * TSHAPS[1][gaussPoint];
    const Real DE2 = 0.5 * TSHAPT[1][gaussPoint];
    const Real DK3 = 0.5 * TSHAPS[2][gaussPoint];
    const Real DE3 = 0.5 * TSHAPT[2][gaussPoint];
    const Real DK4 = 0.5 * TSHAPS[3][gaussPoint];
    const Real DE4 = 0.5 * TSHAPT[3][gaussPoint];

    // Temporary storage arrays
    std::array<Real, 12> sumDerivatives{};    // SU
    std::array<Real, 12> diffDerivatives{};   // DU
    std::vector<Real> ksiDerivatives(6 * numElements);  // DUK
    std::vector<Real> etaDerivatives(6 * numElements);  // DUE

    // Calculate displacement derivatives for each element
    for (int elemIdx = 0; elemIdx < numElements; ++elemIdx) {
        const int offset = 6 * elemIdx;

        // Calculate sum and difference of displacement derivatives
        for (int i = 0; i < 3; ++i) {
            sumDerivatives[0 + i] = displacementDerivatives[elemIdx][4][i] + displacementDerivatives[elemIdx][0][i];
            sumDerivatives[3 + i] = displacementDerivatives[elemIdx][5][i] + displacementDerivatives[elemIdx][1][i];
            sumDerivatives[6 + i] = displacementDerivatives[elemIdx][6][i] + displacementDerivatives[elemIdx][2][i];
            sumDerivatives[9 + i] = displacementDerivatives[elemIdx][7][i] + displacementDerivatives[elemIdx][3][i];

            diffDerivatives[0 + i] = displacementDerivatives[elemIdx][4][i] - displacementDerivatives[elemIdx][0][i];
            diffDerivatives[3 + i] = displacementDerivatives[elemIdx][5][i] - displacementDerivatives[elemIdx][1][i];
            diffDerivatives[6 + i] = displacementDerivatives[elemIdx][6][i] - displacementDerivatives[elemIdx][2][i];
            diffDerivatives[9 + i] = displacementDerivatives[elemIdx][7][i] - displacementDerivatives[elemIdx][3][i];
        }

        // Calculate derivatives with respect to ksi (DK)
        // Constant part
        ksiDerivatives[offset + 0] = DK1 * sumDerivatives[0] + DK2 * sumDerivatives[3] + DK3 * sumDerivatives[6] + DK4 * sumDerivatives[9];
        ksiDerivatives[offset + 1] = DK1 * sumDerivatives[1] + DK2 * sumDerivatives[4] + DK3 * sumDerivatives[7] + DK4 * sumDerivatives[10];
        ksiDerivatives[offset + 2] = DK1 * sumDerivatives[2] + DK2 * sumDerivatives[5] + DK3 * sumDerivatives[8] + DK4 * sumDerivatives[11];
        // Variable part
        ksiDerivatives[offset + 3] = DK1 * diffDerivatives[0] + DK2 * diffDerivatives[3] + DK3 * diffDerivatives[6] + DK4 * diffDerivatives[9];
        ksiDerivatives[offset + 4] = DK1 * diffDerivatives[1] + DK2 * diffDerivatives[4] + DK3 * diffDerivatives[7] + DK4 * diffDerivatives[10];
        ksiDerivatives[offset + 5] = DK1 * diffDerivatives[2] + DK2 * diffDerivatives[5] + DK3 * diffDerivatives[8] + DK4 * diffDerivatives[11];

        // Calculate derivatives with respect to eta (DE)
        // Constant part
        etaDerivatives[offset + 0] = DE1 * sumDerivatives[0] + DE2 * sumDerivatives[3] + DE3 * sumDerivatives[6] + DE4 * sumDerivatives[9];
        etaDerivatives[offset + 1] = DE1 * sumDerivatives[1] + DE2 * sumDerivatives[4] + DE3 * sumDerivatives[7] + DE4 * sumDerivatives[10];
        etaDerivatives[offset + 2] = DE1 * sumDerivatives[2] + DE2 * sumDerivatives[5] + DE3 * sumDerivatives[8] + DE4 * sumDerivatives[11];
        // Variable part
        etaDerivatives[offset + 3] = DE1 * diffDerivatives[0] + DE2 * diffDerivatives[3] + DE3 * diffDerivatives[6] + DE4 * diffDerivatives[9];
        etaDerivatives[offset + 4] = DE1 * diffDerivatives[1] + DE2 * diffDerivatives[4] + DE3 * diffDerivatives[7] + DE4 * diffDerivatives[10];
        etaDerivatives[offset + 5] = DE1 * diffDerivatives[2] + DE2 * diffDerivatives[5] + DE3 * diffDerivatives[8] + DE4 * diffDerivatives[11];
    }

    // Process each node (4 nodes per element)
    for (int nodeIdx = 0; nodeIdx < 4; ++nodeIdx) {
        const int nodeOffset = 3 * nodeIdx;
        const int nodeOffsetHigh = nodeOffset + 12;
        const Real DK = 0.5 * TSHAPS[nodeIdx][gaussPoint];
        const Real DE = 0.5 * TSHAPT[nodeIdx][gaussPoint];

        // Temporary matrices
        std::array<std::array<Real, 3>, 24> bnl0{}; // BNL0
        std::array<std::array<Real, 3>, 24> bnl1{}; // BNL1
        std::array<std::array<Real, 3>, 24> bnl2{}; // BNL2

        for (int elemIdx = 0; elemIdx < numElements; ++elemIdx) {
            const int offset = 6 * elemIdx;
            const int elementType = elementTypes[elemIdx];
            const int thicknessOffset = 4 * (elementType - 1);

            // Fill BNL0 matrix (constant part)
            for (int dim = 0; dim < 3; ++dim) {
                bnl0[nodeOffset + dim][0] = ksiDerivatives[offset + dim] * DK;
                bnl0[nodeOffset + dim][1] = etaDerivatives[offset + dim] * DE;
                bnl0[nodeOffset + dim][2] = ksiDerivatives[offset + dim] * DE + etaDerivatives[offset + dim] * DK;
                
                // Mirror entries for high nodes
                bnl0[nodeOffsetHigh + dim][0] = bnl0[nodeOffset + dim][0];
                bnl0[nodeOffsetHigh + dim][1] = bnl0[nodeOffset + dim][1];
                bnl0[nodeOffsetHigh + dim][2] = bnl0[nodeOffset + dim][2];
            }

            // Fill BNL1 matrix (linear part)
            for (int dim = 0; dim < 3; ++dim) {
                bnl1[nodeOffset + dim][0] = (ksiDerivatives[offset + 3 + dim] - ksiDerivatives[offset + dim]) * DK;
                bnl1[nodeOffset + dim][1] = (etaDerivatives[offset + 3 + dim] - etaDerivatives[offset + dim]) * DE;
                bnl1[nodeOffset + dim][2] = (ksiDerivatives[offset + 3 + dim] - ksiDerivatives[offset + dim]) * DE + 
                                           (etaDerivatives[offset + 3 + dim] - etaDerivatives[offset + dim]) * DK;
                
                // Mirror entries for high nodes
                bnl1[nodeOffsetHigh + dim][0] = (ksiDerivatives[offset + 3 + dim] + ksiDerivatives[offset + dim]) * DK;
                bnl1[nodeOffsetHigh + dim][1] = (etaDerivatives[offset + 3 + dim] + etaDerivatives[offset + dim]) * DE;
                bnl1[nodeOffsetHigh + dim][2] = (ksiDerivatives[offset + 3 + dim] + ksiDerivatives[offset + dim]) * DE + 
                                               (etaDerivatives[offset + 3 + dim] + etaDerivatives[offset + dim]) * DK;
            }

            // Fill BNL2 matrix (quadratic part)
            for (int dim = 0; dim < 3; ++dim) {
                bnl2[nodeOffset + dim][0] = -ksiDerivatives[offset + 3 + dim] * DK;
                bnl2[nodeOffset + dim][1] = -etaDerivatives[offset + 3 + dim] * DE;
                bnl2[nodeOffset + dim][2] = -(ksiDerivatives[offset + 3 + dim] * DE + etaDerivatives[offset + 3 + dim] * DK);
                
                // Mirror entries for high nodes (with sign change)
                bnl2[nodeOffsetHigh + dim][0] = -bnl2[nodeOffset + dim][0];
                bnl2[nodeOffsetHigh + dim][1] = -bnl2[nodeOffset + dim][1];
                bnl2[nodeOffsetHigh + dim][2] = -bnl2[nodeOffset + dim][2];
            }

            // Compute B matrices for each dimension
            for (int dim = 0; dim < 3; ++dim) {
                // Temporary arrays for calculations
                std::array<Real, 6> tmp1{}, tmp2{}, tmp3{}, tmp4{}, utmp{};

                // Compute B1 and B2 matrices for current Gauss point
                for (int i = 0; i < 3; ++i) {
                    tmp1[i] = (1 - Z) * globalMatrix[nodeOffset + dim][thicknessOffset + i][gaussPoint];
                    tmp1[i + 3] = (1 + Z) * globalMatrix[nodeOffset + dim][thicknessOffset + i][gaussPoint];
                    
                    tmp2[i] = bnl0[nodeOffset + dim][i] + Z * bnl1[nodeOffset + dim][i] + Z2 * bnl2[nodeOffset + dim][i];
                    tmp2[i + 3] = bnl0[nodeOffsetHigh + dim][i] + Z * bnl1[nodeOffsetHigh + dim][i] + Z2 * bnl2[nodeOffsetHigh + dim][i];
                    
                    utmp[i] = tmp1[i] + 0.5 * tmp2[i];
                    utmp[i + 3] = tmp1[i + 3] + 0.5 * tmp2[i + 3];
                }

                // Fill B1 matrix
                for (int i = 0; i < 3; ++i) {
                    b1Matrix[elemIdx][nodeOffset + dim][i] = 
                        thicknessCoords[elementType][gaussPoint][0] * utmp[0] +
                        thicknessCoords[elementType][gaussPoint][3] * utmp[1] +
                        thicknessCoords[elementType][gaussPoint][6] * utmp[2];
                    
                    b1Matrix[elemIdx][nodeOffsetHigh + dim][i] = 
                        thicknessCoords[elementType][gaussPoint][0] * utmp[3] +
                        thicknessCoords[elementType][gaussPoint][3] * utmp[4] +
                        thicknessCoords[elementType][gaussPoint][6] * utmp[5];
                }

                // Compute B2 matrix with different combination
                for (int i = 0; i < 3; ++i) {
                    utmp[i] = tmp1[i] + tmp2[i];
                    utmp[i + 3] = tmp1[i + 3] + tmp2[i + 3];
                }

                for (int i = 0; i < 3; ++i) {
                    b2Matrix[elemIdx][nodeOffset + dim][i] = 
                        thicknessCoords[elementType][gaussPoint][0] * utmp[0] +
                        thicknessCoords[elementType][gaussPoint][3] * utmp[1] +
                        thicknessCoords[elementType][gaussPoint][6] * utmp[2];
                    
                    b2Matrix[elemIdx][nodeOffsetHigh + dim][i] = 
                        thicknessCoords[elementType][gaussPoint][0] * utmp[3] +
                        thicknessCoords[elementType][gaussPoint][3] * utmp[4] +
                        thicknessCoords[elementType][gaussPoint][6] * utmp[5];
                }

                // Compute B11 and B12 matrices for high Gauss point
                for (int i = 0; i < 3; ++i) {
                    tmp3[i] = (1 + Z) * globalMatrix[nodeOffset + dim][thicknessOffset + i][gaussPointHigh];
                    tmp3[i + 3] = (1 - Z) * globalMatrix[nodeOffset + dim][thicknessOffset + i][gaussPointHigh];
                    
                    tmp4[i] = bnl0[nodeOffset + dim][i] - Z * bnl1[nodeOffset + dim][i] + Z2 * bnl2[nodeOffset + dim][i];
                    tmp4[i + 3] = bnl0[nodeOffsetHigh + dim][i] - Z * bnl1[nodeOffsetHigh + dim][i] + Z2 * bnl2[nodeOffsetHigh + dim][i];
                    
                    utmp[i] = tmp3[i] + 0.5 * tmp4[i];
                    utmp[i + 3] = tmp3[i + 3] + 0.5 * tmp4[i + 3];
                }

                // Fill B11 matrix
                for (int i = 0; i < 3; ++i) {
                    b11Matrix[elemIdx][nodeOffset + dim][i] = 
                        thicknessCoords[elementType][gaussPointHigh][0] * utmp[0] +
                        thicknessCoords[elementType][gaussPointHigh][3] * utmp[1] +
                        thicknessCoords[elementType][gaussPointHigh][6] * utmp[2];
                    
                    b11Matrix[elemIdx][nodeOffsetHigh + dim][i] = 
                        thicknessCoords[elementType][gaussPointHigh][0] * utmp[3] +
                        thicknessCoords[elementType][gaussPointHigh][3] * utmp[4] +
                        thicknessCoords[elementType][gaussPointHigh][6] * utmp[5];
                }

                // Compute B12 matrix with different combination
                for (int i = 0; i < 3; ++i) {
                    utmp[i] = tmp3[i] + tmp4[i];
                    utmp[i + 3] = tmp3[i + 3] + tmp4[i + 3];
                }

                for (int i = 0; i < 3; ++i) {
                    b12Matrix[elemIdx][nodeOffset + dim][i] = 
                        thicknessCoords[elementType][gaussPointHigh][0] * utmp[0] +
                        thicknessCoords[elementType][gaussPointHigh][3] * utmp[1] +
                        thicknessCoords[elementType][gaussPointHigh][6] * utmp[2];
                    
                    b12Matrix[elemIdx][nodeOffsetHigh + dim][i] = 
                        thicknessCoords[elementType][gaussPointHigh][0] * utmp[3] +
                        thicknessCoords[elementType][gaussPointHigh][3] * utmp[4] +
                        thicknessCoords[elementType][gaussPointHigh][6] * utmp[5];
                }
            }
        }
    }
}